/*
 * Interpolation1DBase has the following inherited classes:
 * - PolInt1D for General Polynomial Interpolation for 1 dim
 * - LinInt1D for Efficient Linear Interpolation for 1 dim
 * - QuaInt1D for Quadratic Interpolation for 1 dim
 * - SplInt1D for Spline Interpolation for 1 dim
 *
 * Author: Yuejia Zhang
 * Date: 2023/2/21
 *
 * */

#ifndef INTERPOLATION_1D
#define INTERPOLATION_1D

#include <cmath>
#include <iostream>

template <class InternalType = double, class ExternalType = float>
class Interpolation1DBase {
 protected:
  int m_;  // number of intervals, one less than number of knots
  ExternalType *xa_, *ya_;
  InternalType *coefficients;

 public:
  /* The construct function computes and stores the coefficients
   * (or second derivatives). xa, ya arrays of x, y. */
  virtual void construct(ExternalType xa[], ExternalType ya[]) = 0;

  /* The evaluate function requires to find the place in the table,
   * at worst O(log N). */
  virtual ExternalType evaluate(ExternalType x) = 0;

  /* subfunction of evaluate, using binary search to locate the interval
   * (xa[i], xa[i+1]) where x is in. returns the left index i.
   * xa must be a sorted array in increasing order!
   * also handles the boundary condition, see the first two lines. */
  int binarySearch(ExternalType *xa, int n, ExternalType x) {
    if (x < xa[0]) return -1;
    if (x > xa[n - 1]) return n;
    int ib = 0, ie = n - 1;
    int im = (ib + ie) / 2;
    while (ib < im) {
      if (x < xa[im])
        ie = im;
      else
        ib = im;
      im = (ib + ie) / 2;
    }
    return im;
  }

  /* subfunction of evaluate, use linear interpolation when x < x_min.
   * the linear function is induced from (xa[ib], xa[ib+1]). */
  ExternalType evaluateLinear(ExternalType x, int ib) {
    ExternalType *xa = this->xa_, *ya = this->ya_;
    InternalType x0 = xa[ib], x1 = xa[ib + 1], y0 = ya[ib], y1 = ya[ib + 1];
    return y0 + (y1 - y0) * (x - x0) / (x1 - x0);
  }

  /* help function for solely debug purpose. */
  virtual void printCoefficients() = 0;
};

/* PolInt1D: Polynomial Interpolation for 1 dim, supports arbitrary N
 * */
template <class InternalType = double, class ExternalType = float, int N = 1>
class PolInt1D : public Interpolation1DBase<InternalType, ExternalType> {
 private:
  ExternalType **xproduct_;
  InternalType *tv_;

 public:
  /* Constructor. Mainly for applying for space.
   * For fixed m and N, the class is constructed only once. */
  PolInt1D(int m) {
    // the number of intervals = the number of points - 1
    m--;
    this->m_ = m;
    // initialize *coefficients, m * (N+1)
    this->coefficients = new InternalType[m * (N + 1)];

    // The following space is used in evaluation
    int NN = N + 1;
    // compute xproduct \prod (x_i-x_j)
    // xproduct[i][j] = (x_i - x_0)...(x_i - x_{j-1})
    xproduct_ = new ExternalType *[NN];
    for (int i = 0; i < NN; i++) xproduct_[i] = new ExternalType[i + 1];
  }

  /* After calling construct, the (N+1) coefficients
   * for each grid (m) are determined, and stored in
   * coefficients[gridIdx : gridIdx + N],
   * where, 0 <= gridIdx < m*(N+1). */
  void construct(ExternalType xa[], ExternalType ya[]) {
    int m = this->m_;
    // initialize **coefficients
    for (int i = 0; i < m * (N + 1); i++) this->coefficients[i] = 0.0;

    this->xa_ = xa;
    this->ya_ = ya;

    // the inner layer
    int ib = (N - 1) / 2;
    int ie = m - N / 2;
    for (int i = ib; i < ie; i++)
      constructGrid(xa, ya, this->coefficients + i * (N + 1), i - ib);

    // the outer layer
    for (int i = ib - 1; i >= 0; i--)
      for (int j = 0; j < (N + 1); j++)
        this->coefficients[i * (N + 1) + j] =
            this->coefficients[(i + 1) * (N + 1) + j];
    for (int i = ie; i < m; i++)
      for (int j = 0; j < (N + 1); j++)
        this->coefficients[i * (N + 1) + j] =
            this->coefficients[(i - 1) * (N + 1) + j];
  }

  /* First find out the place of (x) in grid using binary search,
   * then evaluate the polynomial. */
  ExternalType evaluate(ExternalType x) {
    int ic = this->binarySearch(this->xa_, this->m_ + 1, x);
    if (ic > this->m_ - 1) ic = this->m_ - 1;
    if (ic < 0) {
      return this->evaluateLinear(x, 0);
    }
    // if the point falls in outer layer
    if (ic == 0) ic += (N - 1) / 2;
    if (ic == this->m_ - 1) ic -= N / 2;
    InternalType *cof = this->coefficients + ic * (N + 1);
    int ib = ic - (N - 1) / 2;
    ExternalType sum = 0;
    for (int i = N; i >= 0; i--) {
      sum = sum * (x - this->xa_[ib + i]) + cof[i];
    }
    return sum;
  }

  /* subfunction of construct, compute the coefficients of grid (i, j)
   * and store the array in coefficients[i*n+j].
   * xa[ib:ib+N] and ya[jb:jb+N] are available data to use. */
  void constructGrid(ExternalType xa[], ExternalType ya[], InternalType *ca,
                     int ib) {
    int NN = N + 1;
    // step 1: compute xproduct \prod (x_i-x_j)
    // xproduct[i][j] = (x_i - x_0)...(x_i - x_{j-1})
    ExternalType **xproduct = xproduct_;
    for (int i = 0; i < NN; i++) {
      xproduct[i][0] = 1;
      for (int j = 1; j < i + 1; j++)
        xproduct[i][j] = xproduct[i][j - 1] * (xa[ib + i] - xa[ib + j - 1]);
    }

    // step 2: compute a_{p} by row, 0 <= p < N+1
    InternalType *tv = tv_;
    for (int p = 0; p < NN; p++) {
      InternalType res = ya[ib + p];
      for (int k = 0; k < p; k++) res -= ca[k] * xproduct[p][k];
      res /= xproduct[p][p];
      ca[p] = res;
    }
  }

  void printCoefficients() {
    printf("===Debug: Print Coefficients===\n");
    int m = this->m_;
    for (int i = 0; i < m; i++) {
      printf("Grid (%d) [%.1f, %.1f]: ", i, this->xa_[i], this->xa_[i + 1]);
      for (int k1 = 0; k1 < N + 1; k1++)
        printf("a_{%d} = %.2f; ", k1, this->coefficients[i * (N + 1) + k1]);
      printf("\n");
    }
    printf("================================\n");
  }

  ~PolInt1D() {
    delete[] this->coefficients;
    int NN = N + 1;
    for (int i = 0; i < NN; i++) delete[] xproduct_[i];
    delete[] xproduct_;
  }
};

/* LinInt1D: Efficient Linear Interpolation for 1 dim
 * */
template <class InternalType = double, class ExternalType = float>
class LinInt1D : public Interpolation1DBase<InternalType, ExternalType> {
 public:
  /* Constructor. Mainly for applying for space.
   * For fixed m, the class is constructed only once. */
  LinInt1D(int m) {
    // the number of intervals = the number of points - 1
    m--;
    // initialize *coefficients
    this->coefficients = new InternalType[2 * m];
    this->m_ = m;
  }

  /* The linear form: y = c1 + c2*x.
   * After calling construct, the 2 coefficients for each grid
   * are determined, and stored in coefficients[idx],
   * where 0 <= idx < 2*m. */
  void construct(ExternalType xa[], ExternalType ya[]) {
    int m = this->m_;
    this->xa_ = xa;
    this->ya_ = ya;
    for (int i = 0; i < m; i++)
      constructGrid(xa, ya, this->coefficients + 2 * i, i);
  }

  /* The linear form: y = c1 + c2*x.
   * First find out the place of x in grid using binary search,
   * then evaluate the polynomial. */
  ExternalType evaluate(ExternalType x) {
    int ic = this->binarySearch(this->xa_, this->m_ + 1, x);
    if (ic > this->m_ - 1) ic = this->m_ - 1;
    if (ic < 0) return this->evaluateLinear(x, 0);
    InternalType *cof = this->coefficients + 2 * ic;
    return cof[0] + cof[1] * x;
  }

  /* The linear form: y = c1 + c2*x.
   * subfunction of construct, compute the coefficients of grid i
   * and store the array in coefficients[2*i:2*i+1].*/
  void constructGrid(ExternalType xa[], ExternalType ya[], InternalType *ca,
                     int ib) {
    InternalType c = (ya[ib + 1] - ya[ib]) / (xa[ib + 1] - xa[ib]);
    ca[0] = ya[ib] - c * xa[ib];
    ca[1] = c;
  }

  void printCoefficients() {
    printf("===Debug: Print Coefficients===\n");
    int m = this->m_;
    InternalType *p = this->coefficients;
    for (int i = 0; i < m; i++) {
      printf("Grid (%d) [%.1f, %.1f]: ", i, this->xa_[i], this->xa_[i + 1]);
      printf("c_1 = %.2f, c_2 = %.2f\n", p[0], p[1]);
      p += 2;
    }
    printf("================================\n");
  }

  ~LinInt1D() { delete[] this->coefficients; }
};

/* QuaInt1D: Polynomial Interpolation for 1 dim, N = 2
 * */
template <class InternalType = double, class ExternalType = float>
class QuaInt1D : public Interpolation1DBase<InternalType, ExternalType> {
 public:
  /* Constructor. Mainly for applying for space.
   * For fixed m and N, the class is constructed only once. */
  QuaInt1D(int m) {
    // the number of intervals = the number of points - 1
    m--;
    this->m_ = m;
    // initialize *coefficients, m * (3)
    this->coefficients = new InternalType[m * 3];
  }

  /* After calling construct, the (3) coefficients
   * for each grid (m) are determined, and stored in
   * coefficients[gridIdx : gridIdx + N],
   * where, 0 <= gridIdx < m*(3). */
  void construct(ExternalType xa[], ExternalType ya[]) {
    int m = this->m_;
    // initialize **coefficients
    for (int i = 0; i < m * 3; i++) this->coefficients[i] = 0.0;

    this->xa_ = xa;
    this->ya_ = ya;

    // the inner layer
    int ib = 0;
    int ie = m - 1;
    for (int i = ib; i < ie; i++)
      constructGrid(xa, ya, this->coefficients + i * (3), i - ib);

    // the outer layer
    for (int i = ib - 1; i >= 0; i--)
      for (int j = 0; j < (3); j++)
        this->coefficients[i * (3) + j] = this->coefficients[(i + 1) * (3) + j];
    for (int i = ie; i < m; i++)
      for (int j = 0; j < (3); j++)
        this->coefficients[i * (3) + j] = this->coefficients[(i - 1) * (3) + j];
  }

  /* First find out the place of (x) in grid using binary search,
   * then evaluate the polynomial. */
  ExternalType evaluate(ExternalType x) {
    int ic = this->binarySearch(this->xa_, this->m_ + 1, x);
    if (ic > this->m_ - 1) ic = this->m_ - 1;
    if (ic < 0) {
      return this->evaluateLinear(x, 0);
    }
    // if the point falls in outer layer
    if (ic == this->m_ - 1) ic--;
    InternalType *cof = this->coefficients + ic * (3);
    int ib = ic;
    ExternalType sum = 0;
    for (int i = 2; i >= 0; i--) {
      sum = sum * (x - this->xa_[ib + i]) + cof[i];
    }
    return sum;
  }

  /* subfunction of construct, compute the coefficients of grid (i, j)
   * and store the array in coefficients[i*n+j].
   * xa[ib:ib+N] and ya[jb:jb+N] are available data to use. */
  void constructGrid(ExternalType xa[], ExternalType ya[], InternalType *ca,
                     int ib) {
    int NN = 3;
    // step 1: compute xproduct \prod (x_i-x_j)
    // xproduct[i][j] = (x_i - x_0)...(x_i - x_{j-1})
    ExternalType xproduct[3][3];
    for (int i = 0; i < NN; i++) {
      xproduct[i][0] = 1;
      for (int j = 1; j < i + 1; j++)
        xproduct[i][j] = xproduct[i][j - 1] * (xa[ib + i] - xa[ib + j - 1]);
    }

    // step 2: compute a_{p} by row, 0 <= p < 3
    for (int p = 0; p < NN; p++) {
      InternalType res = ya[ib + p];
      for (int k = 0; k < p; k++) res -= ca[k] * xproduct[p][k];
      res /= xproduct[p][p];
      ca[p] = res;
    }
  }

  void printCoefficients() {
    printf("===Debug: Print Coefficients===\n");
    int m = this->m_;
    for (int i = 0; i < m; i++) {
      printf("Grid (%d) [%.1f, %.1f]: ", i, this->xa_[i], this->xa_[i + 1]);
      for (int k1 = 0; k1 < 3; k1++)
        printf("a_{%d} = %.2f; ", k1, this->coefficients[i * (3) + k1]);
      printf("\n");
    }
    printf("================================\n");
  }

  ~QuaInt1D() { delete[] this->coefficients; }
};

/* SplInt1D: Cubic Spline Interpolation for 1 dim
 * */
template <class InternalType = double, class ExternalType = float>
class SplInt1D : public Interpolation1DBase<InternalType, ExternalType> {
 private:
  InternalType *deriv2, *dx1_, *dx2_, *b_, *T_;

 public:
  /* Constructor. Mainly for applying for space.
   * For fixed m, the class is constructed only once. */
  SplInt1D(int m) {
    this->m_ = m - 1;
    // initialize *coefficients, size = 4 * (m-1)
    this->coefficients = new InternalType[4 * (m - 1)];
    // in construct 1D
    // first-order & second-order difference
    dx1_ = new InternalType[m - 1];
    dx2_ = new InternalType[m - 2];
    b_ = new InternalType[m];
    T_ = new InternalType[m - 2];
    // deriv2 = second derivatives, size = number of knots
    deriv2 = new InternalType[m];
  }

  /* After calling construct, the second-order derivatives M_i
   * for each i are determined, and transformed to coeffcients,
   * which are eventually stored in coefficients[i],
   * where 0 <= i < 4 * m_. */
  void construct(ExternalType xa[], ExternalType ya[]) {
    int n = this->m_ + 1;
    this->xa_ = xa;
    this->ya_ = ya;
    // first-order & second-order difference
    InternalType *dx1 = dx1_;
    InternalType *dx2 = dx2_;
    for (int i = 0; i < n - 1; i++) dx1[i] = xa[i + 1] - xa[i];
    for (int i = 0; i < n - 2; i++) dx2[i] = xa[i + 2] - xa[i];
    InternalType *b = b_;
    for (int i = 0; i < n - 1; i++) b[i] = (ya[i + 1] - ya[i]) / dx1[i];
    for (int i = 0; i < n - 2; i++) b[i] = 6 * (b[i + 1] - b[i]) / dx2[i];

    // solve Tx = b. T stores diagonal of T
    // TODO: Use LAPACK library to solve triangular systems?
    InternalType *T = T_;
    T[0] = 2;
    for (int i = 1; i < n - 2; i++) {
      InternalType scale = dx1[i] / dx2[i] / T[i - 1];
      T[i] = 2 - dx1[i] / dx2[i - 1] * scale;
      b[i] = b[i] - b[i - 1] * scale;
    }
    deriv2[0] = 0;
    deriv2[n - 1] = 0;
    if (n > 2) deriv2[n - 2] = b[n - 3] / T[n - 3];
    for (int i = n - 3; i > 0; i--)
      deriv2[i] = (b[i - 1] - dx1[i] / dx2[i - 1] * deriv2[i + 1]) / T[i - 1];

    // transform deriv2 to coefficients
    for (int i = 0; i < this->m_; i++) {
      int ib = 4 * i;
      InternalType dx = xa[i + 1] - xa[i];
      this->coefficients[ib] = ya[i];
      this->coefficients[ib + 1] =
          (ya[i + 1] - ya[i]) / dx - (deriv2[i + 1] / 6 + deriv2[i] / 3) * dx;
      this->coefficients[ib + 2] = deriv2[i] / 2;
      this->coefficients[ib + 3] = (deriv2[i + 1] - deriv2[i]) / dx / 6;
    }
  }

  /* First find out the place of (x) in grid using binary search,
   * then evaluate the polynomial.
   * f = c0 + c1*x + c2*x*x + c3*x*x*x
   *   = c0 + x*(c1 + x*(c2 + x*c3))
   * */
  ExternalType evaluate(ExternalType x) {
    int ib = this->binarySearch(this->xa_, this->m_ + 1, x);
    if (ib < 0 || ib > this->m_ - 1) {
      // std::cerr << "Spline Interpolation cannot handle extrapolation. "
      //     << "Switch to Linear Interpolation instead. " << std::endl;
      if (ib < 0) ib = 0;
      if (ib > this->m_ - 1) ib = this->m_ - 1;
      return this->evaluateLinear(x, ib);
    }
    int m = this->m_ + 1;
    InternalType *ca = this->coefficients + 4 * ib;
    ExternalType *xa = this->xa_;
    ExternalType *ya = this->ya_;
    x -= xa[ib];
    return ca[0] + x * (ca[1] + x * (ca[2] + x * ca[3]));
  }

  void printCoefficients() {
    printf("===Debug: Print Coefficients===\n");
    int m = this->m_ + 1;
    for (int i = 0; i < m; i++) {
      printf("%.3f\t", this->coefficients[i]);
    }
    printf("================================\n");
  }

  ~SplInt1D() {
    delete[] this->coefficients;
    delete[] dx1_, dx2_, b_, T_;
  }
};

#endif
